import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
import cv2
import numpy as np
import os
from datetime import datetime

class ImageSubscriber(Node):
    def __init__(self):
        super().__init__('image_subscriber')
        self.subscription = self.create_subscription(
            Image,
            'image_topic',
            self.image_callback,
            10)
        self.last_save_time = self.get_clock().now()
        self.save_dir = "saved_images_py"
        os.makedirs(self.save_dir, exist_ok=True)
        self.get_logger().info(f"The image subscriber has started, and the image will be saved to: {self.save_dir}")

    def image_callback(self, msg):
        now = self.get_clock().now()
        if (now - self.last_save_time).nanoseconds / 1e9 < 3.0:
            return
        self.last_save_time = now
        
        try:
            if msg.encoding == 'bgr8':
                dtype = np.uint8
                channels = 3
            else:
                raise ValueError(f"Unsupported image encoding: {msg.encoding}")
            
            image = np.frombuffer(msg.data, dtype=dtype)
            image = image.reshape((msg.height, msg.width, channels))
    
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = os.path.join(self.save_dir, f"image_{timestamp}.png")
            cv2.imwrite(filename, image)
            self.get_logger().info(f"Image saved: {filename}")
            
        except Exception as e:
            self.get_logger().error(f"Image processing error: {str(e)}")

def main(args=None):
    rclpy.init(args=args)
    subscriber = ImageSubscriber()
    rclpy.spin(subscriber)
    subscriber.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()